#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
魔兽争霸3自动对战AI系统 - 可视化界面

该应用提供了系统各模块的可视化操作界面，
包括视觉处理、游戏控制、AI决策等功能。
"""

import sys
import time
import threading
import logging
from datetime import datetime
from typing import Dict, Any, Optional

import cv2
import numpy as np
from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QTabWidget, QWidget, QVBoxLayout, QHBoxLayout,
    QPushButton, QLabel, QTextEdit, QGroupBox, QGridLayout, QLineEdit,
    QCheckBox, QComboBox, QSpinBox, QDoubleSpinBox, QFileDialog, QMessageBox,
    QSplitter, QFrame
)
from PyQt5.QtGui import QPixmap, QImage, QFont, QColor, QPainter, QPen
from PyQt5.QtCore import Qt, QTimer, pyqtSignal

# 导入系统模块
from src.ui.vision_processor import VisionProcessor
from src.game.controller import GameController
from src.ai.decision_maker import DecisionMaker


class LogHandler(logging.Handler):
    """自定义日志处理器，将日志输出到GUI"""
    def __init__(self, text_widget):
        super().__init__()
        self.text_widget = text_widget
        self.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
    
    def emit(self, record):
        msg = self.format(record)
        self.text_widget.append(msg)
        # 自动滚动到底部
        self.text_widget.verticalScrollBar().setValue(
            self.text_widget.verticalScrollBar().maximum()
        )


class GameVisionWidget(QWidget):
    """游戏视觉显示窗口"""
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setMinimumSize(640, 480)
        self.image = None
        
    def set_image(self, image):
        """设置显示的图像"""
        self.image = image
        self.update()
    
    def paintEvent(self, event):
        """绘制图像"""
        super().paintEvent(event)
        if self.image is None:
            return
        
        painter = QPainter(self)
        pixmap = QPixmap.fromImage(self.image)
        # 缩放图像以适应窗口
        scaled_pixmap = pixmap.scaled(
            self.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation
        )
        # 居中显示
        x = (self.width() - scaled_pixmap.width()) // 2
        y = (self.height() - scaled_pixmap.height()) // 2
        painter.drawPixmap(x, y, scaled_pixmap)


class WarcraftAIApp(QMainWindow):
    """魔兽争霸3自动对战AI系统主窗口"""
    
    # 信号定义
    update_log = pyqtSignal(str)
    update_game_state = pyqtSignal(dict)
    
    def __init__(self):
        super().__init__()
        self.setWindowTitle("魔兽争霸3自动对战AI系统")
        self.setGeometry(100, 100, 1200, 800)
        
        # 系统模块实例
        self.vision_processor = None
        self.game_controller = None
        self.decision_maker = None
        
        # 系统状态
        self.is_running = False
        self.is_capturing = False
        self.running_thread = None
        
        # 创建UI
        self.init_ui()
        self.init_logging()
        
    def init_ui(self):
        """初始化用户界面"""
        # 创建中心部件和标签页
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)
        
        # 创建分割器
        splitter = QSplitter(Qt.Vertical)
        main_layout.addWidget(splitter)
        
        # 上部：功能标签页
        tabs = QTabWidget()
        tabs.addTab(self.create_main_tab(), "主控制")
        tabs.addTab(self.create_vision_tab(), "视觉处理")
        tabs.addTab(self.create_control_tab(), "游戏控制")
        tabs.addTab(self.create_ai_tab(), "AI决策")
        tabs.addTab(self.create_settings_tab(), "系统设置")
        splitter.addWidget(tabs)
        
        # 下部：日志显示
        log_group = QGroupBox("系统日志")
        log_layout = QVBoxLayout(log_group)
        self.log_text = QTextEdit()
        self.log_text.setReadOnly(True)
        self.log_text.setLineWrapMode(QTextEdit.WidgetWidth)
        log_layout.addWidget(self.log_text)
        splitter.addWidget(log_group)
        
        # 设置分割器比例
        splitter.setSizes([600, 200])
    
    def create_main_tab(self):
        """创建主控制标签页"""
        tab = QWidget()
        layout = QVBoxLayout(tab)
        
        # 系统状态区域
        status_group = QGroupBox("系统状态")
        status_layout = QGridLayout(status_group)
        
        self.status_label = QLabel("状态: 未初始化")
        self.status_label.setStyleSheet("color: red; font-weight: bold")
        
        self.resources_label = QLabel("资源: 金币: 0, 木材: 0")
        self.units_label = QLabel("单位: 0 个")
        self.buildings_label = QLabel("建筑: 0 个")
        
        status_layout.addWidget(self.status_label, 0, 0, 1, 2)
        status_layout.addWidget(self.resources_label, 1, 0)
        status_layout.addWidget(self.units_label, 1, 1)
        status_layout.addWidget(self.buildings_label, 2, 0)
        
        # 控制按钮区域
        control_group = QGroupBox("系统控制")
        control_layout = QHBoxLayout(control_group)
        
        self.init_button = QPushButton("初始化系统")
        self.start_button = QPushButton("开始运行")
        self.stop_button = QPushButton("停止运行")
        self.stop_button.setEnabled(False)
        
        self.init_button.clicked.connect(self.initialize_system)
        self.start_button.clicked.connect(self.start_system)
        self.stop_button.clicked.connect(self.stop_system)
        
        control_layout.addWidget(self.init_button)
        control_layout.addWidget(self.start_button)
        control_layout.addWidget(self.stop_button)
        
        # 游戏视觉显示
        vision_group = QGroupBox("游戏视觉")
        vision_layout = QVBoxLayout(vision_group)
        self.vision_widget = GameVisionWidget()
        vision_layout.addWidget(self.vision_widget)
        
        # 添加到主布局
        layout.addWidget(status_group)
        layout.addWidget(control_group)
        layout.addWidget(vision_group, 1)  # 1表示伸展因子
        
        return tab
    
    def create_vision_tab(self):
        """创建视觉处理标签页"""
        tab = QWidget()
        layout = QVBoxLayout(tab)
        
        # 视觉设置
        settings_group = QGroupBox("视觉设置")
        settings_layout = QGridLayout(settings_group)
        
        settings_layout.addWidget(QLabel("屏幕区域:"), 0, 0)
        
        self.region_x = QSpinBox()
        self.region_x.setRange(0, 3840)
        self.region_x.setValue(0)
        self.region_y = QSpinBox()
        self.region_y.setRange(0, 2160)
        self.region_y.setValue(0)
        self.region_width = QSpinBox()
        self.region_width.setRange(100, 3840)
        self.region_width.setValue(1920)
        self.region_height = QSpinBox()
        self.region_height.setRange(100, 2160)
        self.region_height.setValue(1080)
        
        region_hbox = QHBoxLayout()
        region_hbox.addWidget(QLabel("X:"))
        region_hbox.addWidget(self.region_x)
        region_hbox.addWidget(QLabel("Y:"))
        region_hbox.addWidget(self.region_y)
        region_hbox.addWidget(QLabel("宽度:"))
        region_hbox.addWidget(self.region_width)
        region_hbox.addWidget(QLabel("高度:"))
        region_hbox.addWidget(self.region_height)
        
        settings_layout.addLayout(region_hbox, 0, 1)
        
        # 模板加载
        template_group = QGroupBox("模板管理")
        template_layout = QVBoxLayout(template_group)
        
        template_hbox = QHBoxLayout()
        self.template_name = QLineEdit("建筑模板")
        self.load_template_button = QPushButton("加载模板")
        self.load_template_button.clicked.connect(self.load_template)
        
        template_hbox.addWidget(QLabel("模板名称:"))
        template_hbox.addWidget(self.template_name)
        template_hbox.addWidget(self.load_template_button)
        
        self.template_list = QTextEdit()
        self.template_list.setReadOnly(True)
        self.template_list.setMinimumHeight(100)
        
        template_layout.addLayout(template_hbox)
        template_layout.addWidget(self.template_list)
        
        # 测试按钮
        test_group = QGroupBox("视觉测试")
        test_layout = QHBoxLayout(test_group)
        
        self.capture_button = QPushButton("捕获屏幕")
        self.detect_resources_button = QPushButton("检测资源")
        self.analyze_state_button = QPushButton("分析游戏状态")
        
        self.capture_button.clicked.connect(self.capture_screen)
        self.detect_resources_button.clicked.connect(self.detect_resources)
        self.analyze_state_button.clicked.connect(self.analyze_game_state)
        
        test_layout.addWidget(self.capture_button)
        test_layout.addWidget(self.detect_resources_button)
        test_layout.addWidget(self.analyze_state_button)
        
        # 结果显示
        result_group = QGroupBox("视觉结果")
        result_layout = QVBoxLayout(result_group)
        
        self.vision_result = QTextEdit()
        self.vision_result.setReadOnly(True)
        result_layout.addWidget(self.vision_result)
        
        # 添加到主布局
        layout.addWidget(settings_group)
        layout.addWidget(template_group)
        layout.addWidget(test_group)
        layout.addWidget(result_group, 1)
        
        return tab
    
    def create_control_tab(self):
        """创建游戏控制标签页"""
        tab = QWidget()
        layout = QVBoxLayout(tab)
        
        # 控制设置
        settings_group = QGroupBox("控制设置")
        settings_layout = QGridLayout(settings_group)
        
        settings_layout.addWidget(QLabel("游戏窗口区域:"), 0, 0)
        
        self.window_x = QSpinBox()
        self.window_x.setRange(0, 3840)
        self.window_x.setValue(0)
        self.window_y = QSpinBox()
        self.window_y.setRange(0, 2160)
        self.window_y.setValue(0)
        self.window_width = QSpinBox()
        self.window_width.setRange(100, 3840)
        self.window_width.setValue(1920)
        self.window_height = QSpinBox()
        self.window_height.setRange(100, 2160)
        self.window_height.setValue(1080)
        
        window_hbox = QHBoxLayout()
        window_hbox.addWidget(QLabel("X:"))
        window_hbox.addWidget(self.window_x)
        window_hbox.addWidget(QLabel("Y:"))
        window_hbox.addWidget(self.window_y)
        window_hbox.addWidget(QLabel("宽度:"))
        window_hbox.addWidget(self.window_width)
        window_hbox.addWidget(QLabel("高度:"))
        window_hbox.addWidget(self.window_height)
        
        settings_layout.addLayout(window_hbox, 0, 1)
        
        # 鼠标控制
        mouse_group = QGroupBox("鼠标控制")
        mouse_layout = QGridLayout(mouse_group)
        
        mouse_layout.addWidget(QLabel("X坐标:"), 0, 0)
        mouse_layout.addWidget(QLabel("Y坐标:"), 1, 0)
        
        self.mouse_x = QSpinBox()
        self.mouse_x.setRange(0, 3840)
        self.mouse_x.setValue(500)
        self.mouse_y = QSpinBox()
        self.mouse_y.setRange(0, 2160)
        self.mouse_y.setValue(500)
        
        mouse_layout.addWidget(self.mouse_x, 0, 1)
        mouse_layout.addWidget(self.mouse_y, 1, 1)
        
        mouse_buttons = QHBoxLayout()
        self.move_mouse_button = QPushButton("移动鼠标")
        self.left_click_button = QPushButton("左键点击")
        self.right_click_button = QPushButton("右键点击")
        
        self.move_mouse_button.clicked.connect(self.move_mouse)
        self.left_click_button.clicked.connect(self.left_click)
        self.right_click_button.clicked.connect(self.right_click)
        
        mouse_buttons.addWidget(self.move_mouse_button)
        mouse_buttons.addWidget(self.left_click_button)
        mouse_buttons.addWidget(self.right_click_button)
        
        mouse_layout.addLayout(mouse_buttons, 2, 0, 1, 2)
        
        # 键盘控制
        keyboard_group = QGroupBox("键盘控制")
        keyboard_layout = QGridLayout(keyboard_group)
        
        keyboard_layout.addWidget(QLabel("按键:"), 0, 0)
        self.key_input = QLineEdit("a")
        keyboard_layout.addWidget(self.key_input, 0, 1)
        
        keyboard_buttons = QHBoxLayout()
        self.press_key_button = QPushButton("按键")
        self.press_key_button.clicked.connect(self.press_key)
        keyboard_buttons.addWidget(self.press_key_button)
        
        keyboard_layout.addLayout(keyboard_buttons, 1, 0, 1, 2)
        
        # 游戏操作
        game_group = QGroupBox("游戏操作")
        game_layout = QGridLayout(game_group)
        
        game_layout.addWidget(QLabel("建筑/单位快捷键:"), 0, 0)
        self.game_hotkey = QLineEdit("b")
        game_layout.addWidget(self.game_hotkey, 0, 1)
        
        game_buttons = QHBoxLayout()
        self.build_button = QPushButton("建造建筑")
        self.train_button = QPushButton("训练单位")
        self.attack_move_button = QPushButton("攻击移动")
        
        self.build_button.clicked.connect(self.build_structure)
        self.train_button.clicked.connect(self.train_unit)
        self.attack_move_button.clicked.connect(self.attack_move)
        
        game_buttons.addWidget(self.build_button)
        game_buttons.addWidget(self.train_button)
        game_buttons.addWidget(self.attack_move_button)
        
        game_layout.addLayout(game_buttons, 1, 0, 1, 2)
        
        # 添加到主布局
        layout.addWidget(settings_group)
        layout.addWidget(mouse_group)
        layout.addWidget(keyboard_group)
        layout.addWidget(game_group)
        
        return tab
    
    def create_ai_tab(self):
        """创建AI决策标签页"""
        tab = QWidget()
        layout = QVBoxLayout(tab)
        
        # AI设置
        ai_group = QGroupBox("AI设置")
        ai_layout = QGridLayout(ai_group)
        
        # 模型配置
        ai_layout.addWidget(QLabel("模型文件:"), 0, 0)
        self.model_path = QLineEdit()
        self.browse_model_button = QPushButton("浏览...")
        self.browse_model_button.clicked.connect(self.browse_model)
        ai_layout.addWidget(self.model_path, 0, 1)
        ai_layout.addWidget(self.browse_model_button, 0, 2)
        
        # 大模型设置
        self.use_llm_checkbox = QCheckBox("使用大语言模型辅助决策")
        self.use_llm_checkbox.stateChanged.connect(self.toggle_llm_settings)
        ai_layout.addWidget(self.use_llm_checkbox, 1, 0, 1, 3)
        
        ai_layout.addWidget(QLabel("API地址:"), 2, 0)
        self.llm_api_url = QLineEdit("http://localhost:8000/v1/completions")
        ai_layout.addWidget(self.llm_api_url, 2, 1, 1, 2)
        
        ai_layout.addWidget(QLabel("API密钥:"), 3, 0)
        self.llm_api_key = QLineEdit()
        ai_layout.addWidget(self.llm_api_key, 3, 1, 1, 2)
        
        # RL参数调整
        rl_group = QGroupBox("强化学习参数")
        rl_layout = QGridLayout(rl_group)
        
        rl_layout.addWidget(QLabel("Gamma (折扣因子):"), 0, 0)
        self.gamma = QDoubleSpinBox()
        self.gamma.setRange(0.0, 1.0)
        self.gamma.setValue(0.95)
        self.gamma.setSingleStep(0.01)
        rl_layout.addWidget(self.gamma, 0, 1)
        
        rl_layout.addWidget(QLabel("Epsilon (探索率):"), 1, 0)
        self.epsilon = QDoubleSpinBox()
        self.epsilon.setRange(0.0, 1.0)
        self.epsilon.setValue(1.0)
        self.epsilon.setSingleStep(0.01)
        rl_layout.addWidget(self.epsilon, 1, 1)
        
        rl_layout.addWidget(QLabel("学习率:"), 2, 0)
        self.learning_rate = QDoubleSpinBox()
        self.learning_rate.setRange(0.0001, 0.1)
        self.learning_rate.setValue(0.001)
        self.learning_rate.setSingleStep(0.0001)
        self.learning_rate.setDecimals(4)
        rl_layout.addWidget(self.learning_rate, 2, 1)
        
        # 模型操作
        model_group = QGroupBox("模型操作")
        model_layout = QHBoxLayout(model_group)
        
        self.load_model_button = QPushButton("加载模型")
        self.save_model_button = QPushButton("保存模型")
        self.update_params_button = QPushButton("更新参数")
        
        self.load_model_button.clicked.connect(self.load_ai_model)
        self.save_model_button.clicked.connect(self.save_ai_model)
        self.update_params_button.clicked.connect(self.update_ai_params)
        
        model_layout.addWidget(self.load_model_button)
        model_layout.addWidget(self.save_model_button)
        model_layout.addWidget(self.update_params_button)
        
        # AI状态
        state_group = QGroupBox("AI状态")
        state_layout = QVBoxLayout(state_group)
        
        self.ai_state = QTextEdit()
        self.ai_state.setReadOnly(True)
        state_layout.addWidget(self.ai_state)
        
        # 添加到主布局
        layout.addWidget(ai_group)
        layout.addWidget(rl_group)
        layout.addWidget(model_group)
        layout.addWidget(state_group, 1)
        
        # 初始禁用大模型设置
        self.toggle_llm_settings()
        
        return tab
    
    def create_settings_tab(self):
        """创建设置标签页"""
        tab = QWidget()
        layout = QVBoxLayout(tab)
        
        # 系统设置
        system_group = QGroupBox("系统设置")
        system_layout = QGridLayout(system_group)
        
        system_layout.addWidget(QLabel("循环间隔 (秒):"), 0, 0)
        self.loop_interval = QDoubleSpinBox()
        self.loop_interval.setRange(0.1, 10.0)
        self.loop_interval.setValue(2.0)
        self.loop_interval.setSingleStep(0.1)
        system_layout.addWidget(self.loop_interval, 0, 1)
        
        system_layout.addWidget(QLabel("日志级别:"), 1, 0)
        self.log_level = QComboBox()
        self.log_level.addItems(["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"])
        self.log_level.setCurrentText("INFO")
        self.log_level.currentTextChanged.connect(self.change_log_level)
        system_layout.addWidget(self.log_level, 1, 1)
        
        # 关于信息
        about_group = QGroupBox("关于")
        about_layout = QVBoxLayout(about_group)
        
        about_text = QTextEdit()
        about_text.setReadOnly(True)
        about_text.setHtml("""
        <h3>魔兽争霸3自动对战AI系统</h3>
        <p>版本: 1.0.0</p>
        <p>一个基于视觉识别和强化学习的魔兽争霸3自动对战系统</p>
        <p>功能模块:</p>
        <ul>
            <li>视觉处理: 捕获和分析游戏画面</li>
            <li>游戏控制: 发送鼠标和键盘指令</li>
            <li>AI决策: 使用强化学习制定策略</li>
        </ul>
        """)
        
        about_layout.addWidget(about_text)
        
        # 添加到主布局
        layout.addWidget(system_group)
        layout.addWidget(about_group, 1)
        
        return tab
    
    def init_logging(self):
        """初始化日志系统"""
        logger = logging.getLogger()
        logger.setLevel(logging.INFO)
        
        # 清除现有处理器
        for handler in logger.handlers[:]:
            logger.removeHandler(handler)
        
        # 添加GUI日志处理器
        gui_handler = LogHandler(self.log_text)
        logger.addHandler(gui_handler)
        
        # 添加文件日志处理器
        file_handler = logging.FileHandler('data/logs/warcraft_ai.log', encoding='utf-8')
        file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
        logger.addHandler(file_handler)
    
    def change_log_level(self, level):
        """更改日志级别"""
        level_map = {
            "DEBUG": logging.DEBUG,
            "INFO": logging.INFO,
            "WARNING": logging.WARNING,
            "ERROR": logging.ERROR,
            "CRITICAL": logging.CRITICAL
        }
        logger = logging.getLogger()
        logger.setLevel(level_map.get(level, logging.INFO))
        logging.info(f"日志级别已更改为: {level}")
    
    def initialize_system(self):
        """初始化系统模块"""
        try:
            logging.info("开始初始化系统模块...")
            
            # 初始化视觉处理模块
            screen_region = (
                self.region_x.value(),
                self.region_y.value(),
                self.region_width.value(),
                self.region_height.value()
            )
            self.vision_processor = VisionProcessor(screen_region)
            logging.info(f"视觉处理模块初始化完成，区域: {screen_region}")
            
            # 初始化游戏控制模块
            game_window = (
                self.window_x.value(),
                self.window_y.value(),
                self.window_width.value(),
                self.window_height.value()
            )
            self.game_controller = GameController(game_window)
            logging.info(f"游戏控制模块初始化完成，窗口: {game_window}")
            
            # 初始化AI决策模块
            use_llm = self.use_llm_checkbox.isChecked()
            llm_api_url = self.llm_api_url.text() if use_llm else None
            llm_api_key = self.llm_api_key.text() if use_llm else None
            
            model_path = self.model_path.text() if self.model_path.text() else None
            self.decision_maker = DecisionMaker(
                model_path=model_path,
                use_llm=use_llm,
                llm_api_url=llm_api_url,
                llm_api_key=llm_api_key
            )
            
            # 更新AI参数
            self.update_ai_params_internal()
            
            logging.info("AI决策模块初始化完成")
            logging.info("系统初始化成功！")
            
            # 更新状态
            self.status_label.setText("状态: 已初始化")
            self.status_label.setStyleSheet("color: green; font-weight: bold")
            
            # 启用开始按钮
            self.start_button.setEnabled(True)
            self.init_button.setEnabled(False)
            
        except Exception as e:
            logging.error(f"系统初始化失败: {e}")
            QMessageBox.critical(self, "初始化失败", f"系统初始化失败: {str(e)}")
    
    def start_system(self):
        """启动自动对战系统"""
        if not self.vision_processor or not self.game_controller or not self.decision_maker:
            QMessageBox.warning(self, "错误", "请先初始化系统！")
            return
        
        try:
            self.is_running = True
            self.running_thread = threading.Thread(target=self.main_loop)
            self.running_thread.daemon = True
            self.running_thread.start()
            
            logging.info("魔兽争霸3自动对战AI系统已启动")
            
            # 更新UI状态
            self.start_button.setEnabled(False)
            self.stop_button.setEnabled(True)
            self.status_label.setText("状态: 运行中")
            self.status_label.setStyleSheet("color: blue; font-weight: bold")
            
        except Exception as e:
            logging.error(f"启动系统失败: {e}")
            self.is_running = False
            QMessageBox.critical(self, "启动失败", f"启动系统失败: {str(e)}")
    
    def stop_system(self):
        """停止自动对战系统"""
        try:
            self.is_running = False
            if self.running_thread:
                self.running_thread.join(timeout=3.0)
            
            logging.info("魔兽争霸3自动对战AI系统已停止")
            
            # 更新UI状态
            self.start_button.setEnabled(True)
            self.stop_button.setEnabled(False)
            self.status_label.setText("状态: 已初始化")
            self.status_label.setStyleSheet("color: green; font-weight: bold")
            
        except Exception as e:
            logging.error(f"停止系统失败: {e}")
            QMessageBox.critical(self, "停止失败", f"停止系统失败: {str(e)}")
    
    def main_loop(self):
        """主循环"""
        while self.is_running:
            try:
                # 1. 捕获并分析游戏画面
                logging.info("捕获游戏画面...")
                game_state = self.vision_processor.analyze_game_state()
                logging.info(f"当前游戏状态: {game_state}")
                
                # 更新UI显示
                self.update_ui_with_game_state(game_state)
                
                # 2. AI制定决策
                logging.info("AI制定决策...")
                decision = self.decision_maker.make_decision(game_state)
                logging.info(f"AI决策: {decision}")
                
                # 更新AI状态显示
                self.update_ai_state_display(decision)
                
                # 3. 执行操作
                if decision.get('action'):
                    self.execute_decision(decision)
                
                # 控制循环频率
                time.sleep(self.loop_interval.value())
                
            except Exception as e:
                logging.error(f"运行出错: {e}")
                time.sleep(1.0)  # 出错后等待1秒继续
    
    def update_ui_with_game_state(self, game_state):
        """使用游戏状态更新UI"""
        # 更新资源显示
        resources = game_state.get('resources', {})
        self.resources_label.setText(f"资源: 金币: {resources.get('gold', 0)}, 木材: {resources.get('wood', 0)}")
        
        # 更新单位和建筑显示
        unit_count = game_state.get('unit_count', 0)
        building_count = game_state.get('building_count', 0)
        self.units_label.setText(f"单位: {unit_count} 个")
        self.buildings_label.setText(f"建筑: {building_count} 个")
        
        # 更新视觉显示（如果有截图）
        frame = self.vision_processor.capture_screen()
        if frame is not None:
            # 转换OpenCV图像到Qt格式
            rgb_image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            h, w, ch = rgb_image.shape
            bytes_per_line = ch * w
            q_image = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
            self.vision_widget.set_image(q_image)
    
    def update_ai_state_display(self, decision):
        """更新AI状态显示"""
        state_text = f"决策来源: {decision.get('source', 'unknown')}\n"
        state_text += f"动作类型: {decision.get('action', {}).get('type', 'idle')}\n"
        state_text += f"动作参数: {decision.get('action', {}).get('params', {})}\n"
        
        if 'reason' in decision:
            state_text += f"决策原因: {decision['reason']}\n"
        
        if 'confidence' in decision:
            state_text += f"置信度: {decision['confidence']:.2f}\n"
        
        # 在主线程中更新UI
        from PyQt5.QtCore import QEvent, QCoreApplication
        
        def update_ui():
            self.ai_state.setPlainText(state_text)
        
        QCoreApplication.postEvent(self, QEvent(QEvent.User))
        QTimer.singleShot(0, update_ui)
    
    def execute_decision(self, decision):
        """执行AI决策"""
        action = decision.get('action', {})
        action_type = action.get('type', 'idle')
        params = action.get('params', {})
        
        try:
            if action_type == 'move' and params:
                x = params.get('x', 0)
                y = params.get('y', 0)
                logging.info(f"执行移动操作到 ({x}, {y})")
                self.game_controller.move_mouse(x, y)
                self.game_controller.click(x, y)
            
            elif action_type == 'build' and params:
                structure = params.get('structure', '')
                x = params.get('x', 0)
                y = params.get('y', 0)
                logging.info(f"执行建造操作: {structure} 在 ({x}, {y})")
                # 这里可以根据建筑类型映射到对应的快捷键
                hotkey = self.game_hotkey.text()
                self.game_controller.build_structure(hotkey, x, y)
            
            elif action_type == 'train' and params:
                unit = params.get('unit', '')
                logging.info(f"执行训练操作: {unit}")
                # 这里可以根据单位类型映射到对应的快捷键
                hotkey = self.game_hotkey.text()
                self.game_controller.train_unit(hotkey)
            
            elif action_type == 'attack' and params:
                x = params.get('x', 0)
                y = params.get('y', 0)
                logging.info(f"执行攻击操作到 ({x}, {y})")
                self.game_controller.attack_move(x, y)
            
        except Exception as e:
            logging.error(f"执行决策失败: {e}")
    
    # 视觉处理相关方法
    def capture_screen(self):
        """捕获屏幕"""
        if not self.vision_processor:
            QMessageBox.warning(self, "错误", "视觉处理模块未初始化！")
            return
        
        try:
            frame = self.vision_processor.capture_screen()
            if frame is not None:
                # 转换OpenCV图像到Qt格式
                rgb_image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                h, w, ch = rgb_image.shape
                bytes_per_line = ch * w
                q_image = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
                self.vision_widget.set_image(q_image)
                self.vision_result.setPlainText(f"捕获成功，图像尺寸: {w}x{h}")
            else:
                self.vision_result.setPlainText("捕获失败")
        except Exception as e:
            self.vision_result.setPlainText(f"捕获异常: {str(e)}")
    
    def detect_resources(self):
        """检测资源"""
        if not self.vision_processor:
            QMessageBox.warning(self, "错误", "视觉处理模块未初始化！")
            return
        
        try:
            frame = self.vision_processor.capture_screen()
            if frame is not None:
                resources = self.vision_processor.detect_resources(frame)
                self.vision_result.setPlainText(f"资源检测结果:\n金币: {resources.get('gold', 0)}\n木材: {resources.get('wood', 0)}")
            else:
                self.vision_result.setPlainText("无法捕获屏幕进行资源检测")
        except Exception as e:
            self.vision_result.setPlainText(f"资源检测异常: {str(e)}")
    
    def analyze_game_state(self):
        """分析游戏状态"""
        if not self.vision_processor:
            QMessageBox.warning(self, "错误", "视觉处理模块未初始化！")
            return
        
        try:
            game_state = self.vision_processor.analyze_game_state()
            state_text = f"游戏状态分析结果:\n"
            state_text += f"时间戳: {datetime.fromtimestamp(game_state.get('timestamp', time.time())).strftime('%Y-%m-%d %H:%M:%S')}\n"
            state_text += f"资源: {game_state.get('resources', {})}\n"
            state_text += f"小地图: {'已提取' if game_state.get('minimap') is not None else '未找到'}\n"
            state_text += f"单位数量: {len(game_state.get('units', []))}\n"
            state_text += f"建筑数量: {len(game_state.get('buildings', []))}\n"
            
            self.vision_result.setPlainText(state_text)
            # 同时更新UI显示
            self.update_ui_with_game_state(game_state)
        except Exception as e:
            self.vision_result.setPlainText(f"游戏状态分析异常: {str(e)}")
    
    def load_template(self):
        """加载图像模板"""
        if not self.vision_processor:
            QMessageBox.warning(self, "错误", "视觉处理模块未初始化！")
            return
        
        try:
            template_name = self.template_name.text()
            if not template_name:
                QMessageBox.warning(self, "错误", "请输入模板名称！")
                return
            
            file_path, _ = QFileDialog.getOpenFileName(
                self, "选择模板图像", "", "图像文件 (*.png *.jpg *.jpeg *.bmp)"
            )
            
            if file_path:
                success = self.vision_processor.load_template(template_name, file_path)
                if success:
                    self.template_list.append(f"✓ {template_name}: {file_path}")
                    self.vision_result.setPlainText(f"模板加载成功: {template_name}")
                else:
                    self.vision_result.setPlainText(f"模板加载失败: {template_name}")
        except Exception as e:
            self.vision_result.setPlainText(f"模板加载异常: {str(e)}")
    
    # 游戏控制相关方法
    def move_mouse(self):
        """移动鼠标"""
        if not self.game_controller:
            QMessageBox.warning(self, "错误", "游戏控制模块未初始化！")
            return
        
        try:
            x = self.mouse_x.value()
            y = self.mouse_y.value()
            self.game_controller.move_mouse(x, y)
            logging.info(f"手动移动鼠标到 ({x}, {y})")
        except Exception as e:
            logging.error(f"移动鼠标失败: {e}")
    
    def left_click(self):
        """左键点击"""
        if not self.game_controller:
            QMessageBox.warning(self, "错误", "游戏控制模块未初始化！")
            return
        
        try:
            x = self.mouse_x.value()
            y = self.mouse_y.value()
            self.game_controller.click(x, y, button='left')
            logging.info(f"左键点击位置 ({x}, {y})")
        except Exception as e:
            logging.error(f"左键点击失败: {e}")
    
    def right_click(self):
        """右键点击"""
        if not self.game_controller:
            QMessageBox.warning(self, "错误", "游戏控制模块未初始化！")
            return
        
        try:
            x = self.mouse_x.value()
            y = self.mouse_y.value()
            self.game_controller.click(x, y, button='right')
            logging.info(f"右键点击位置 ({x}, {y})")
        except Exception as e:
            logging.error(f"右键点击失败: {e}")
    
    def press_key(self):
        """按键"""
        if not self.game_controller:
            QMessageBox.warning(self, "错误", "游戏控制模块未初始化！")
            return
        
        try:
            key = self.key_input.text()
            if key:
                self.game_controller.press_key(key)
                logging.info(f"按下键: {key}")
            else:
                QMessageBox.warning(self, "错误", "请输入要按下的键！")
        except Exception as e:
            logging.error(f"按键失败: {e}")
    
    def build_structure(self):
        """建造建筑"""
        if not self.game_controller:
            QMessageBox.warning(self, "错误", "游戏控制模块未初始化！")
            return
        
        try:
            hotkey = self.game_hotkey.text()
            x = self.mouse_x.value()
            y = self.mouse_y.value()
            if hotkey:
                self.game_controller.build_structure(hotkey, x, y)
                logging.info(f"建造建筑，快捷键: {hotkey}，位置: ({x}, {y})")
            else:
                QMessageBox.warning(self, "错误", "请输入建筑快捷键！")
        except Exception as e:
            logging.error(f"建造建筑失败: {e}")
    
    def train_unit(self):
        """训练单位"""
        if not self.game_controller:
            QMessageBox.warning(self, "错误", "游戏控制模块未初始化！")
            return
        
        try:
            hotkey = self.game_hotkey.text()
            if hotkey:
                self.game_controller.train_unit(hotkey)
                logging.info(f"训练单位，快捷键: {hotkey}")
            else:
                QMessageBox.warning(self, "错误", "请输入单位快捷键！")
        except Exception as e:
            logging.error(f"训练单位失败: {e}")
    
    def attack_move(self):
        """攻击移动"""
        if not self.game_controller:
            QMessageBox.warning(self, "错误", "游戏控制模块未初始化！")
            return
        
        try:
            x = self.mouse_x.value()
            y = self.mouse_y.value()
            self.game_controller.attack_move(x, y)
            logging.info(f"攻击移动到位置 ({x}, {y})")
        except Exception as e:
            logging.error(f"攻击移动失败: {e}")
    
    # AI相关方法
    def browse_model(self):
        """浏览模型文件"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择模型文件", "", "模型文件 (*.h5 *.keras *.weights)"
            )
        if file_path:
            self.model_path.setText(file_path)
    
    def load_ai_model(self):
        """加载AI模型"""
        if not self.decision_maker:
            QMessageBox.warning(self, "错误", "AI决策模块未初始化！")
            return
        
        try:
            model_path = self.model_path.text()
            if not model_path:
                model_path, _ = QFileDialog.getOpenFileName(
                    self, "选择模型文件", "", "模型文件 (*.h5 *.keras *.weights)"
                )
                if not model_path:
                    return
                self.model_path.setText(model_path)
            
            self.decision_maker.load_model(model_path)
            logging.info(f"模型加载成功: {model_path}")
            QMessageBox.information(self, "成功", "模型加载成功！")
        except Exception as e:
            logging.error(f"模型加载失败: {e}")
            QMessageBox.critical(self, "错误", f"模型加载失败: {str(e)}")
    
    def save_ai_model(self):
        """保存AI模型"""
        if not self.decision_maker:
            QMessageBox.warning(self, "错误", "AI决策模块未初始化！")
            return
        
        try:
            model_path, _ = QFileDialog.getSaveFileName(
                self, "保存模型文件", "data/models/model.weights", "模型文件 (*.weights)"
            )
            
            if model_path:
                self.decision_maker.save_model(model_path)
                logging.info(f"模型保存成功: {model_path}")
                QMessageBox.information(self, "成功", "模型保存成功！")
        except Exception as e:
            logging.error(f"模型保存失败: {e}")
            QMessageBox.critical(self, "错误", f"模型保存失败: {str(e)}")
    
    def update_ai_params(self):
        """更新AI参数"""
        if not self.decision_maker:
            QMessageBox.warning(self, "错误", "AI决策模块未初始化！")
            return
        
        try:
            self.update_ai_params_internal()
            logging.info("AI参数更新成功")
            QMessageBox.information(self, "成功", "AI参数更新成功！")
        except Exception as e:
            logging.error(f"AI参数更新失败: {e}")
            QMessageBox.critical(self, "错误", f"AI参数更新失败: {str(e)}")
    
    def update_ai_params_internal(self):
        """内部更新AI参数"""
        if not self.decision_maker:
            return
        
        # 更新强化学习参数
        self.decision_maker.gamma = self.gamma.value()
        self.decision_maker.epsilon = self.epsilon.value()
        self.decision_maker.learning_rate = self.learning_rate.value()
    
    def toggle_llm_settings(self):
        """切换大模型设置的启用状态"""
        enabled = self.use_llm_checkbox.isChecked()
        self.llm_api_url.setEnabled(enabled)
        self.llm_api_key.setEnabled(enabled)
    
    def closeEvent(self, event):
        """窗口关闭事件"""
        if self.is_running:
            reply = QMessageBox.question(
                self, '确认退出', '系统正在运行中，确定要退出吗？',
                QMessageBox.Yes | QMessageBox.No, QMessageBox.No
            )
            if reply != QMessageBox.Yes:
                event.ignore()
                return
            
            # 停止系统
            self.stop_system()
        
        event.accept()


def main():
    """主函数"""
    app = QApplication(sys.argv)
    app.setApplicationName("魔兽争霸3自动对战AI系统")
    window = WarcraftAIApp()
    window.show()
    sys.exit(app.exec_())

if __name__ == "__main__":
    main()